import os
import subprocess
from multiprocessing import Process, Event
import signal
import atexit
import sys
import time
import argparse


def run_server(stop_event, server_name, host, port, cpu, ram, rom):
    def signal_handler(signum, frame):
        stop_event.set()

    # 注册信号处理函数
    signal.signal(signal.SIGTERM, signal_handler)

    # 设置环境变量
    env = os.environ.copy()
    env['VIM_SERVER_NAME'] = server_name
    env['VIM_RESOURCE_CPU'] = cpu
    env['VIM_RESOURCE_RAM'] = ram
    env['VIM_RESOURCE_ROM'] = rom

    # 构建uvicorn命令
    command = ['uvicorn','run:app',
        '--host', host,
        '--port', port,
        '--log-level', 'error',
    ]

    with subprocess.Popen(command, env=env) as proc:
        try:
            while not stop_event.is_set():
                time.sleep(0.5)
        except Exception as e:
            print(f"Exception in server process: {e}")
        finally:
            proc.terminate()
            proc.wait()  # 等待子进程结束


def run_clean():
    command = ['python','clean.py','-db']
    with subprocess.Popen(command) as proc:
        proc.wait()


def create_vim_process(stop_event, **kwargs):
    vim_name = kwargs.get('vim_name','vim_default')
    vim_ip = kwargs.get('vim_ip','127.0.0.1')
    vim_port = kwargs.get('vim_port','5000')
    vim_cpu = kwargs.get('vim_cpu','4')
    vim_ram = kwargs.get('vim_ram','1')
    vim_rom = kwargs.get('vim_rom','32')
    vimProcess = Process(target=run_server, args=(stop_event, vim_name, vim_ip, vim_port, 
                                                  vim_cpu, vim_ram, vim_rom))
    vimProcess.daemon = True
    vimProcess.start()

    return vimProcess


def create_clean_process(vim_processes):
    for p in vim_processes:
        p.terminate()
        p.join()
    run_clean()


def parse_args():
    parser = argparse.ArgumentParser(description="VIM (VIM Server) configuration")
    parser.add_argument('--vim_name', required=True, help='The name of the VIM (e.g. vim_1)')
    parser.add_argument('--vim_ip', required=True, help='The IP address of the VIM (e.g. 192.168.31.79)')
    parser.add_argument('--vim_port', required=True, help='The port number of the VIM (e.g. 5001)')
    parser.add_argument('--vim_cpu', required=True, help='The CPU/core requirement of the VIM (e.g. 4)')
    parser.add_argument('--vim_ram', required=True, help='The RAM/GB requirement of the VIM (e.g. 2)')
    parser.add_argument('--vim_rom', required=True, help='The ROM/GB requirement of the VIM (e.g. 128)')

    args = parser.parse_args()

    return vars(args)


if __name__ == '__main__':

    vim_info = parse_args()

    stop_event = Event()
    vim_process = create_vim_process(stop_event, **vim_info)
    vim_processes = [vim_process]
    atexit.register(create_clean_process, vim_processes=vim_processes)

    try:
        while True:
            timestamp = time.time()
            local_time = time.localtime(timestamp)
            formatted_time = time.strftime("%H:%M:%S", local_time)
            print(f"{formatted_time} {vim_info['vnf_name']} {vim_info['vim_ip']}:{vim_info['vim_port']}")
            time.sleep(5)
    except KeyboardInterrupt:
        print('Main process received interrupt, initiating clean up...')
        stop_event.set()
        create_clean_process(vim_processes)
        print('Cleanup completed, exiting...')
        sys.exit(0)

